import matplotlib.pylab as plt
import norm
import scipy.stats as ss
import pandas as pd
import statsmodels.api as sm
import numpy as np
from tqdm import tqdm
import Judge


# 参考
# https://blog.csdn.net/weixin_43230383/article/details/121607538

# 检测残差,回归异常值,残差正态检验,残差正态画图
# regression_result:回归结果,可以不填
# regression_dict:用于回归的字典形如{y:ydata,x1:xdata}这样,必须填
# regression_str:用于回归的字符串形如{'y~x1'}里面的这个字符串这样,必须填
# title:图片标题
# 返回正态检验的结果矩阵
def check(regression_result='回归结果,可以不填',
          regression_dict='用于回归的字典形如{y:ydata,x1:xdata}这样,必须填',
          regression_str='用于回归的字符串形如{\'y~x1\'}里面的这个字符串这样,必须填',
          title='标题'):
    # 如果回归模型为空
    if regression_result == '回归结果,可以不填':
        if regression_dict == '用于回归的字典形如{y:ydata,x1:xdata}这样,必须填' or regression_str == '用于回归的字符串形如{\'y~x1\'}里面的这个字符串这样,必须填':
            print('没有回归结果,请输入完整数据')
        else:
            # 重新回归
            regression_result = sm.formula.ols(regression_str, regression_dict).fit()
    outerindexs = []
    # 寻找异常值之并绘制图
    resid, outer_result, upp, low, outer_index, r, num = DrawErrorBar(regression_result, regression_dict, title=title,
                                                                      Draw=False)
    print('\n异常值数量:' + str(np.sum(outer_result['bonf(p)'] != 1)), end='')
    outerindexs.append(outer_index)
    # 预测
    # 记录剔除的次数
    Delete_times = 0
    # 记录有异常值时的回归结果,用于组合画图
    Nums = []
    # 这个是用于画图的
    outer_draw_index = []
    resids = []
    R = []

    # 循环剔除异常值
    while np.sum(outer_result['bonf(p)'] != 1) > 0:
        Delete_times += 1
        # dataframe
        outer_index = outer_result[outer_result['bonf(p)'] < 1].index.values
        # 绘制图
        resid, outer_result, upp, low, outer_index, r, num = DrawErrorBar(regression_result, regression_dict,
                                                                          title=title, Draw=False)
        # 记录数据
        Nums.append(num)
        outer_draw_index.append(outer_index)
        resids.append(resid)
        R.append(r)
        # 去除异常值
        for key, item in regression_dict.items():
            item = np.array(item)
            item = np.delete(item, outer_index)
            # 更新值
            regression_dict[key] = item
        # 重新回归
        regression_result = sm.formula.ols(regression_str, regression_dict).fit()
        outerindexs.append(outer_index)
        # print(outer_result)  # 输出已知数据的野值检验
        outer_result = regression_result.outlier_test()
        print('->' + str(np.sum(outer_result['bonf(p)'] != 1)), end='')
    # 输出删除的异常值的索引
    # print('\n异常值的索引:')
    # for i in range(len(outerindexs)):
    #     print('第' + str(i + 1) + '回归的异常值的索引:\n', outerindexs[i])
    # print('残差的方差', regression_result.mse_resid)
    popvalue = None
    # 不画删除最后一个正常的
    # if len(Nums)>=1:
    # popvalue=Nums.pop(-1)
    # 绘制误差棒组合图
    # 绘制进度条
    if len(Nums) == 4:
        fig, axes = plt.subplots(2, 2, sharex=True, sharey=True)
        # 一次画多张图
        # 2x2图
        print('画图程序:')
        for i in tqdm(range(len(Nums))):
            num = Nums[i]
            r = R[i]
            resid = resids[i]
            outer_index = outer_draw_index[i]
            ROW_2x2 = np.ceil((i + 1) / 2)
            ROW_2x2 = int(ROW_2x2) - 1
            Column = Judge.Zero2Top(i + 1, 2)
            Column -= 1
            axes[ROW_2x2, Column].errorbar(num, resid, r, fmt='o-')
            axes[ROW_2x2, Column].errorbar(num[outer_index], resid[outer_index], r[outer_index], fmt='or')
            # axes[ROW_2x2, Column].title.set_text(title + '残差分布')
    # 画3列图,或一行图
    elif len(Nums) != 0:
        # 行数

        # 是否只能画一行图
        ROW_1xn = (len(Nums) <= 3)
        if ROW_1xn:
            # 特殊情况,画异常图和原图
            if len(Nums) == 1:
                # Nums.append(popvalue)
                # fig, axes = plt.subplots(1, len(Nums), sharex=True, sharey=True)
                for i in tqdm(range(len(Nums))):
                    num = Nums[i]
                    r = R[i]
                    resid = resids[i]
                    outer_index = outer_draw_index[i]
                    # axes[i].errorbar(num, resid, r, fmt='o-')
                    # axes[i].errorbar(num[outer_index], resid[outer_index], r[outer_index], fmt='or')
                    plt.suptitle(title + '残差分布')
                    plt.errorbar(Nums[0], resids[0], R[0], fmt='o-')
                    plt.errorbar(Nums[0][outer_draw_index[0]], resids[0][outer_draw_index[0]],
                                 R[0][outer_draw_index[0]], fmt='or')
            else:
                fig, axes = plt.subplots(1, len(Nums), sharex=True, sharey=True)
                for i in tqdm(range(len(Nums))):
                    num = Nums[i]
                    r = R[i]
                    resid = resids[i]
                    outer_index = outer_draw_index[i]
                    axes[i].errorbar(num, resid, r, fmt='o-')
                    axes[i].errorbar(num[outer_index], resid[outer_index], r[outer_index], fmt='or')
        else:
            ROW = np.ceil(len(Nums) / 3)
            ROW = int(ROW)
            fig, axes = plt.subplots(ROW, 3, sharex=True, sharey=True)
            for i in tqdm(range(len(Nums))):
                num = Nums[i]
                r = R[i]
                resid = resids[i]
                outer_index = outer_draw_index[i]

                # 获得行索引
                ROW_3x3 = np.ceil((i + 1) / 3)
                ROW_3x3 = int(ROW_3x3) - 1
                # 获得列索引
                Column = Judge.Zero2Top(i + 1, 3)
                Column -= 1
                axes[ROW_3x3, Column].errorbar(num, resid, r, fmt='o-')
                axes[ROW_3x3, Column].errorbar(num[outer_index], resid[outer_index], r[outer_index], fmt='or')

        plt.subplots_adjust(wspace=0, hspace=0)
        # 不画大标题更好看
        # plt.suptitle(title + '残差分布')

    # 绘制删除异常值之后的图
    resid, outer_result, upp, low, outer_index, r, num = DrawErrorBar(regression_result, regression_dict,
                                                                      title=title, Draw=True)
    # 统计结果
    print(regression_result.summary())
    # 绘制拟合图
    # 回归系数表汇总提取
    coef_df = None
    coef_series = None
    if 1 > 0:  # 1>0用于收折
        coef_df = pd.DataFrame({"回归系数": regression_result.params,  # 回归系数
                                "回归系数标准差": regression_result.bse,  # 回归系数标准差
                                # "回归系数T值": round(regression_result.tvalues, 3),  # 回归系数T值
                                # "p-values": round(regression_result.pvalues, 3)  # 回归系数P值
                                })
        # 回归系数置信区间 默认5%，括号中可填具体数字 比如0.05, 0.1
        coef_df[['置信区间0.025', '置信区间_0.975']] = regression_result.conf_int()
        coef_series = pd.Series({
            "R方": regression_result.rsquared,
            "调整R方": regression_result.rsquared_adj,
            "AIC": regression_result.aic,
            "BIC": regression_result.bic,
            "F统计量": regression_result.fvalue,
            "p": regression_result.f_pvalue,
            "模型mse": regression_result.mse_model,
            "残差mse": regression_result.mse_resid,
            "总体mse": regression_result.mse_total,
            "模型SSR": regression_result.ssr,
            "最大似然值": regression_result.llf,
        })

    # 获取预测值
    Predict_Value = regression_result.predict(regression_dict)
    plt.figure()
    plt.plot(Predict_Value, label='预测值')
    plt.plot(regression_dict['y'], label='初始值')
    plt.plot(upp, label='置信上限')
    plt.plot(low, label='置信下限')
    plt.legend()
    plt.title(title + '拟合图')

    # 正态检验的结果
    Normal_Result = norm.Normal_Inspection(resid)
    # 绘制正态图
    norm.draw(title, data=resid)
    # 正态检验的结果

    return regression_result, Normal_Result, coef_df, coef_series


def DrawErrorBar(regression_result, regression_dict, title, Draw=True):
    pre = regression_result.get_prediction(regression_dict)
    df = pre.summary_frame(alpha=0.05)
    dfv = df.values
    low, upp = dfv[:, 4:].T  # 置信下限上限
    r = (upp - low) / 2  # 置信半径
    num = np.arange(1, len(regression_dict['y']) + 1)
    # 获得残差
    resid = regression_result.resid
    resid = np.array(resid)

    outer_result = regression_result.outlier_test()
    # print(outer_result)  # 输出已知数据的野值检验

    # 找出异常值
    # 绘制误差区间图-误差棒状图
    # 获得异常误差的索引
    outer_index = outer_result[outer_result['bonf(p)'] != 1].index.values
    # print('异常值的索引:\n', outer_index)
    if Draw:
        plt.figure()
        # errorbar参数
        # https://baijiahao.baidu.com/s?id=1737342038816904018&wfr=spider&for=pc
        # http://www.qb5200.com/article/234525.html
        plt.errorbar(num, resid, r, fmt='o-', capsize=3.5, capthick=1.5)
        plt.errorbar(num[outer_index], resid[outer_index], r[outer_index], fmt='or')
        plt.title(title + '残差分布')
    return resid, outer_result, upp, low, outer_index, r, num
